
import matplotlib.pyplot as plt
import pandas as pd
from pev3.data.DataModule import DataModule
from pev3.factor.FactorModule import FactorModule
from pev3.util.DataBase import DB_CONN
from pev3.util.StockUtil import getTradingDates

"""
零投资组合
"""
class ZeroInvestmentPortfolioAnalysis:
    """
    零投资组合的初始化方法

    :param factor:  因子名字
    :param startDate: 分析的开始日期
    :param endDate: 分析的结束日期
    :param interval: 调整周期，交易日数
    :param position: 档位数，默认划分为10档
    :param ascending: 是否按照因子值正序排列，默认为正序
    """
    def __init__(self, factor, startDate, endDate, interval, position=10, ascending=True):
        # 单期收益的DataFrame
        self.dfProfit = pd.DataFrame(
            columns=['top', 'bottom', 'portfolio']
        )
        # 累计收益的DataFrame,沪深300作为基准
        self.cumulativeProfit = pd.DataFrame(
            columns=['top', 'bottom', 'portfolio', 'hs300']
        )
        # 单期股次数的DataFrame
        self.dfCount = pd.DataFrame(
            columns=['top','bottom']
        )

        # 净值
        self.lastTopNetValue = 1
        self.lastBottomNetValue = 1
        self.lastPortfolioNetValue = 1
        # 沪深300首日
        self.hs300FirstValue = -1
        # 因子名字
        self.factor = factor
        # 分析的日期范围
        self.startDate = startDate
        self.endDate = endDate
        # 调整周期
        self.interval = interval
        # 排序方式
        self.ascending = ascending
        # 档位数
        self.position = position

    """
    计算某一档的平均收益
    :param dfDailies: 日行情的DataFrame
    :param positionCodeCloseDict:
    :return: 收益
    """
    def computAverageProfit(self, dfDailies, positionCodeCloseDict):
        #提取股票列表
        positionCodes = list(positionCodeCloseDict.keys())

        # 只有存在股票时，才进行计算
        if len(positionCodes) > 0:
            # 所有股票代码
            codes = set(dfDailies.index)

            # 所有股票的累计收益
            profitSum = 0
            # 实际参与统计的股票数
            count = 0
            # 计算所有股票的收益
            for code in positionCodes:
                count += 1
                buyClose = positionCodeCloseDict[code]

                # 计算所有股票的累计收益
                if code in codes:
                    profitSum += (dfDailies.loc[code]['close'] - buyClose) / buyClose

            # 计算单期平均收益
            return round(profitSum * 100/ count, 2), count

        #没有数据时，返回None
        return None


    """
    计算收益
    :param lastAdjustDate: 上一个调整日
    :param dfDailies:
    :param topDailies:
    :param bottomDailies:
    :param adjustDate: 当前调整日
    :return:
    """
    def computeProfit(self, lastAdjustDate, dfDailies, topDailies, bottomDailies, adjustDate):
        # 只有存在上一个调整日，才计算上期的收益
        if lastAdjustDate is not None:
            # 计算首档收益
            topProfit = self.computAverageProfit(dfDailies, topDailies)

            # 计算末档收益
            bottomProfit = self.computAverageProfit(dfDailies, bottomDailies)

            # 计算组合收益
            portfolioProfit = topProfit[0] - bottomProfit[0]

            # 添加结果到DataFrame中
            self.dfProfit.loc[lastAdjustDate] = {
                'top': topProfit[0],
                'bottom': bottomProfit[0],
                'portfolio': portfolioProfit
            }

            # 计算累积收益（复利方式)
            # 首档
            topCumulativeProfit = round((self.lastTopNetValue * (1 + topProfit[0] / 100) - 1) * 100, 2)
            self.lastTopNetValue *= (1 + topProfit[0] / 100)
            #末档
            bottomCumulativeProfit = round((self.lastBottomNetValue * (1 + bottomProfit[0] / 100) - 1) * 100, 2)
            self.lastBottomNetValue *= (1 + bottomProfit[0] / 100)

            # 组合
            portfolioCumulativeProfit = round((self.lastPortfolioNetValue * ( 1 + portfolioProfit / 100) - 1) * 100, 2)
            self.lastPortfolioNetValue *= (1 + portfolioProfit / 100)

            # 计算沪深300的累计收益
            dm = DataModule()
            hs300K = dm.getKData('000300', index=True, startDate=adjustDate, endDate=adjustDate)
            hs300K.set_index(['date'], 1, inplace=True)
            hs300Profit = (hs300K.loc[adjustDate]['close'] - self.hs300FirstValue) / self.hs300FirstValue

            self.cumulativeProfit.loc[lastAdjustDate] = {
                'top': topCumulativeProfit,
                'bottom': bottomCumulativeProfit,
                'portfolio': portfolioCumulativeProfit,
                'hs300': hs300Profit
            }

            self.dfCount.loc[lastAdjustDate] = {
                'top': topProfit[1],
                'bottom': bottomProfit[1]
            }

    """
    市值的零投资组合分析
    """
    def analyze(self):
        # 初始化对数据管理子系统接口的调用
        dm = DataModule()
        # 初始化对因子管理子系统接口的调用
        fm = FactorModule()

        # 获取分析周期内的
        allDates = getTradingDates(self.startDate, self.endDate)

        # 首档和末档，股票代码和后复权价格的Dictionary
        topDailies = dict()
        bottomDailies = dict()

        # 暂存上一个调整
        lastAdjustDate = None

        # 设置沪深300的首日值
        hs300K = dm.getKData('000300', index=True, startDate=allDates[0], endDate=allDates[0])
        self.hs300FirstValue = hs300K.loc[0]['close']

        # 计算每日收益
        for index in range(0, len(allDates), self.interval):
            adjustDate = allDates[index]

            # 获取因子值，按照指定的顺序排序
            dfFactor = fm.getSingleDateFactors(self.factor, adjustDate)
            if dfFactor.index.size == 0:
                continue
            dfFactor.sort_values(self.factor, ascending=self.ascending, inplace=True)
            #将股票代码设为index
            dfFactor.set_index(['code'], inplace=True)

            # 获取当日所有股票的行情
            dfDailies = dm.getOneDayKData(autype='hfq', date=adjustDate)
            # 将code设为index
            dfDailies.set_index(['code'], inplace=True)

            # 计算收益
            self.computeProfit(lastAdjustDate, dfDailies, topDailies, bottomDailies, adjustDate)


